# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from coin.exchange.base.tick import FixedTickPrice
from coin.strategy.mm.subscription import FeedSubscriptionRequest

import experimental.prophet as prophet
from coin.experimental.jaewon.research.adaptive_edge import features, lt_model
from experimental.prophet.ops.linear_model import linear_regression_predict
from experimental.prophet.craps.universe import universe_btc


def _model(symbol_nick,
           ref_symbol_nicks,
           linear_model,
           date,
           sampler_offsets=None,
           lt_model_param=None,
           debug=False):
  sampler_offsets = sampler_offsets or {}
  lt_model_param = lt_model_param or {}

  # Time
  ts = prophet.global_timestamp()

  # Symbols
  universe = {elem.nickname: elem for elem in universe_btc(date)}
  symbol_info = universe[symbol_nick]
  feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)

  book = prophet.get_book_from_fastfeed(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs
  ask0_p = book.estimated_ask0_price.named_ref('ask0_p')
  bid0_p = book.estimated_bid0_price.named_ref('bid0_p')

  # Features
  with prophet.namespace('feature'):
    feature_vars = features.features(ref_symbol_nicks + [symbol_nick], date)

  # Sampler
  with prophet.namespace('sampler'):
    sampler_vars = features.samplers('feature/')
    sampler_vars = {v.name: v for v in sampler_vars}

  # Trigger
  trigger = features.trigger(symbol_nick, 'feature/')

  # Long-term model theo price
  with prophet.namespace('lt_model'):
    theo_ask0_p, theo_bid0_p = lt_model.theo_price(symbol_info, date=date, **lt_model_param)

  # Model
  with prophet.namespace('edge_model'):
    x_vars = [v for v in feature_vars if not v.name.endswith('.x')]

    buy_edge_bp = 15.
    sell_edge_bp = 15.

    for sampler_name in linear_model.keys():
      offset = sampler_offsets.get(sampler_name, 3.)
      # logging.info('Sampler %s: %f bp', sampler_name, offset)

      sampler = sampler_vars[sampler_name]
      buy_model = dict(linear_model[sampler_name]['buy']['model'].params)
      sell_model = dict(linear_model[sampler_name]['sell']['model'].params)
      with prophet.control_if(sampler):
        sampler_buy_edge_bp = linear_regression_predict(x_vars, buy_model) + float(offset)
        sampler_sell_edge_bp = linear_regression_predict(x_vars, sell_model) + float(offset)

      buy_edge_bp = prophet.cond(sampler, sampler_buy_edge_bp, buy_edge_bp)
      sell_edge_bp = prophet.cond(sampler, sampler_sell_edge_bp, sell_edge_bp)

    buy_p = (theo_bid0_p * (1. - buy_edge_bp / 10000.)).named_ref('buy_p')
    sell_p = (theo_ask0_p * (1. + sell_edge_bp / 10000.)).named_ref('sell_p')

  if debug:
    prophet.printer('%t buy_edge_bp=%s sell_edge_bp=%s', ts, buy_edge_bp, sell_edge_bp)


class Calculator:
  def __init__(self, tick=1., machine='', eval_callback=None, **kwargs):
    self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick

    self._graph = prophet.Graph()
    with prophet.as_default_graph(self._graph):
      _model(**kwargs)

    self._ask0_p_var = self._graph.get_named_variable('ask0_p')
    self._bid0_p_var = self._graph.get_named_variable('bid0_p')
    self._buy_p_var = self._graph.get_named_variable('edge_model/buy_p')
    self._sell_p_var = self._graph.get_named_variable('edge_model/sell_p')
    self._timestamp = self._graph.get_named_variable('global_timestamp')
    self._trigger = self._graph.get_named_variable('trigger')

    self._adapter = prophet.FastfeedGraphAdapter(self._graph, machine, eval_callback=eval_callback)

  @property
  def ready(self):
    return not self._buy_p_var.is_null and not self._sell_p_var.is_null

  @property
  def trigger(self):
    return self._trigger.value

  @property
  def timestamp(self):
    return self._timestamp.value

  def calculate_price(self, force_passive=False):
    assert self.ready

    if force_passive:
      ask0_minus_one = self._tick.get_prev_price(self._ask0_p_var.value)
      bid0_plus_one = self._tick.get_next_price(self._bid0_p_var.value)
      buy_price = min(self._tick.floor(self._buy_p_var.value), ask0_minus_one, bid0_plus_one)
      sell_price = max(self._tick.ceil(self._sell_p_var.value), ask0_minus_one, bid0_plus_one)
    else:
      buy_price = self._tick.floor(self._buy_p_var.value)
      sell_price = self._tick.ceil(self._sell_p_var.value)

    return buy_price, sell_price

  def get_feed_sub_reqs(self, feed_sub_req=None):
    feed_sub_req = feed_sub_req or FeedSubscriptionRequest()
    return self._adapter.get_feed_sub_requests(feed_sub_req)

  @property
  def book_reset_callbacks(self):
    return [self._adapter.on_book_reset]
